#include "nominallowrankextract.h"

#include <ctype.h>
#include <errno.h>
#include <libgen.h>
#include <math.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>

static float
qfunc_data (Data* data,
            float hq,
            unsigned int n_threads)
{
  return qfunc (n_threads,
                data->loga,
                data->n_workers,
                data->rank,
                data->gamma,
                data->logbeta,
                data->n_images,
                data->c,
                data->n_labels,
                data->d,
                data->pz,
                data->ratings,
                data->n_ratings,
                &data->priorloga.base,
                &data->priorgamma.base,
                &data->priorlogbeta.base,
                &data->priorc.base,
                &data->priord.base,
                hq);
}

static float
estep_data (Data*        data,
            unsigned int n_threads)
{
  return estep (n_threads,
                data->loga,
                data->rank,
                data->logbeta,
                data->n_images,
                data->c,
                data->n_labels,
                data->d,
                data->pz,
                data->priorz,
                data->clamp,
                data->ratings,
                data->n_ratings);
}

static void
mstep_data (Data*               data,
            unsigned int        n_passes,
            float               eta)
{
  return mstep_sgd (data->loga,
                    data->n_workers,
                    data->rank,
                    data->gamma,
                    data->logbeta,
                    data->n_images,
                    data->c,
                    data->n_labels,
                    data->d,
                    data->pz,
                    data->ratings,
                    data->n_ratings,
                    data->image_ratings,
                    data->worker_ratings,
                    &data->priorloga.base,
                    &data->priorgamma.base,
                    &data->priorlogbeta,
                    &data->priorc.base,
                    &data->priord.base,
                    n_passes,
                    eta);
}

#define maybe_fprintf(...)      \
  if (! quiet)                  \
    {                           \
      fprintf (__VA_ARGS__);    \
    }

static void
em (Data*  data,
    float  eta,
    float  decay,
    unsigned int n_passes,
    unsigned int n_threads,
    int    quiet)
{
  float q;
  float lastq;
  const float epsilon = 1e-4;
  unsigned int badstep = 0;
  Data save;

  memset (&save, 0, sizeof (save));

  do
    {
      int whoops;

      maybe_fprintf (stderr, "estep start");
      float hq = estep_data (data, n_threads);
      maybe_fprintf (stderr, " end");
      q = qfunc_data (data, hq, n_threads);
      lastq = q;
      data_cpy (&save, data);

      do
        {
          maybe_fprintf (stderr, " mstep start: q = %g", lastq);
          mstep_data (data, n_passes, eta);
          q = qfunc_data (data, hq, n_threads);
          maybe_fprintf (stderr, " mstep end: q = %g\n", q);

          whoops = (lastq > q || fabs ((q - lastq) / lastq) < epsilon);

          if (whoops)
            {
              ++badstep;
              eta *= decay;
              data_cpy (data, &save);
            }
        }
      while (whoops && badstep < 5);
    }
  while (badstep < 5);

  data_destruct (&save);
}

static void
print_loga (FILE*       fp,
            const Data* data)
{
  for (unsigned int i = 0; i < data->n_workers; ++i)
    {
      fprintf (fp, "a[%u] =", i);

      for (unsigned int j = 0; j < data->rank; ++j)
        {
          fprintf (fp, " %f", exp (data->loga[i * data->rank + j]));
        }

      fprintf (fp, "\n");
    }
}

static void
print_logbeta (FILE*       fp,
               const Data* data)
{
  for (unsigned int i = 0; i < data->n_images; ++i)
    {
      fprintf (fp, "Beta[%u] = %f\n", i, exp (data->logbeta[i]));
    }
}

static void
print_gamma (FILE*       fp,
             const Data* data)
{
  fprintf (fp, "Gamma =");

  for (unsigned int j = 0; j < data->rank; ++j)
    {
      fprintf (fp, " %f", data->gamma[j]);
    }

  fprintf (fp, "\n");
}

static void
print_c (FILE*       fp,
         const Data* data)
{
  fprintf (fp, "C =");

  for (unsigned int i = 0; i < data->rank; ++i)
    {
      for (unsigned int j = 0; j < data->n_labels; ++j)
        {
          fprintf (fp, " %f", data->c[i * data->n_labels + j]);
        }
    }

  fprintf (fp, "\n");
}

static void
print_d (FILE*       fp,
         const Data* data)
{
  fprintf (fp, "D =");

  for (unsigned int i = 0; i < data->rank; ++i)
    {
      for (unsigned int j = 0; j < data->n_labels; ++j)
        {
          fprintf (fp, " %f", data->d[i * data->n_labels + j]);
        }
    }

  fprintf (fp, "\n");
}

static void
print_pz (FILE*         fp,
          const Data*   data)
{
  for (unsigned int i = 0; i < data->n_images; ++i)
    {
      fprintf (fp, "P(Z(%u)=(k>0)) =", i);

      for (unsigned int j = 1; j < data->n_labels; ++j)
        {
          fprintf (fp, " %f", data->pz[j + i * data->n_labels]);
        }

      fprintf (fp, "\n");
    }
}

static void
print_parameters (FILE*         fp,
                  const Data*   data)
{
  print_loga (fp, data);
  print_logbeta (fp, data);
  print_gamma (fp, data);
  print_c (fp, data);
  print_d (fp, data);
  print_pz (fp, data);
}

static void
usage (char* progname)
{
  fprintf (stderr, 
           "usage: %s -r rank [--] [-h] [-d decay] [-l eta] [-p passes] [-q] [-s seed] [-t threads] [datafile] [clampfile]\n"
           "\n"
           "  -r: specify approximation rank (required).\n"
           "\n"
           "  datafile: specify file to read input data from.\n"
           "      if datafile eq '-' use standard input\n"
           "      if not specified, \n"
           "        if standard input is not a tty, use standard input\n"
           "        if standard input is a tty, display this message.\n"
           "\n"
           "  clampfile: specify file to read clamp specification from.\n"
           "      if clampfile eq '-' use standard input\n"
           "        if datafile is also standard input, datafile is read first\n"
           "      if not specified, \n"
           "        use same file as datafile; datafile is read first\n"
           "\n"
           "  --: interpret subsequent -whatever as a file argument\n"
           "\n"
           "  -h: display this message.\n"
           "\n"
           "  -d: specify learning rate decay after bad mstep (default: 0.707).\n"
           "\n"
           "  -l: specify learning rate (default: 1).\n"
           "\n"
           "  -p: specify data passes per m-step (default: 10).\n"
           "\n"
           "  -q: quiet (don't display progress).\n"
           "\n"
           "  -s: specify random seed (default: 45).\n"
           "\n"
           "  -t: specify number of threads (default: 1).\n"
           "\n",
           basename (progname)
          );

  exit (1);
}

static int
parse_positive_float (const char*      bufptr,
                      float*           valptr,
                      float            maxval,
                      const char*      vname)
{
  char* endptr;

  *valptr = strtod (bufptr, &endptr);

  if (*endptr == '\0')
    {
      if (*valptr > maxval)
        {
          fprintf (stderr, "error reading %s: value %f exceeding maximum value %f\n",
                   vname,
                   *valptr,
                   maxval);

          return -1;
        }

      return 0;
    }

  fprintf (stderr, "error reading %s: bad character '%c' (%hhx)\n",
           vname,
           *endptr,
           *endptr);

  return -1;
}

static int
parse_unsigned_int (const char*      bufptr,
                    unsigned int*    valptr,
                    unsigned int     maxval,
                    const char*      vname)
{
  char* endptr;

  *valptr = strtoul (bufptr, &endptr, 10);

  if (*endptr == '\0')
    {
      if (*valptr >= maxval)
        {
          fprintf (stderr, "error reading %s: value %u at or exceeding maximum value %u\n",
                   vname,
                   *valptr,
                   maxval);

          return -1;
        }

      return 0;
    }

  fprintf (stderr, "error reading %s: bad character '%c' (%hhx)\n",
           vname,
           *endptr,
           *endptr);

  return -1;
}

int 
main (int       argc, 
      char*     argv[])
{
  Data data;
  float eta = 0;
  float decay = 0;
  unsigned int n_passes = 0;
  unsigned int n_threads = 0;
  unsigned int rank = 0;
  unsigned int seed = 0;
  int quiet = 0;
  char* progname = argv[0];

  while (argc > 1)
    {
      if (strcmp (argv[1], "--") == 0)
        {
          ++argv;
          --argc;
          break;
        }
      else if (strcmp (argv[1], "-h") == 0)
        {
          usage (progname);
        }
      else if (! quiet && strcmp (argv[1], "-q") == 0)
        {
          quiet = 1;
          ++argv;
          --argc;
        }
      else if (eta == 0 && strcmp (argv[1], "-l") == 0 && argc > 2)
        {
          ++argv;
          --argc;
          if (parse_positive_float (argv[1], &eta, 1000, "eta"))
            {
              exit (1);
            }
          ++argv;
          --argc;
        }
      else if (decay == 0 && strcmp (argv[1], "-d") == 0 && argc > 2)
        {
          ++argv;
          --argc;
          if (parse_positive_float (argv[1], &decay, 1, "decay"))
            {
              exit (1);
            }
          ++argv;
          --argc;
        }
      else if (! rank && strcmp (argv[1], "-r") == 0 && argc > 2)
        {
          ++argv;
          --argc;
          if (parse_unsigned_int (argv[1], &rank, 1000, "rank"))
            {
              exit (1);
            }
          ++argv;
          --argc;
        }
      else if (! n_passes && strcmp (argv[1], "-p") == 0 && argc > 2)
        {
          ++argv;
          --argc;
          if (parse_unsigned_int (argv[1], &n_passes, 1000, "n_passes"))
            {
              exit (1);
            }
          ++argv;
          --argc;
        }
      else if (! n_passes && strcmp (argv[1], "-t") == 0 && argc > 2)
        {
          ++argv;
          --argc;
          if (parse_unsigned_int (argv[1], &n_threads, 1000, "n_threads"))
            {
              exit (1);
            }
          ++argv;
          --argc;
        }
      else if (! seed && strcmp (argv[1], "-s") == 0 && argc > 2)
        {
          ++argv;
          --argc;
          if (parse_unsigned_int (argv[1], &seed, 0xFFFFFFFFUL, "seed"))
            {
              exit (1);
            }
          ++argv;
          --argc;
        }
      else
        {
          break;
        }
    }

  if (argc > 3 || (argc == 1 && isatty (0)))
    {
      usage (progname);
    }

  if (rank == 0)
    {
      fprintf (stderr, "fatal: rank not specified\n");
      exit (1);
    }

  maybe_fprintf (stderr, "using rank of %u\n", rank);

  if (seed == 0)
    {
      seed = 45;
    }

  maybe_fprintf (stderr, "using random seed of %u\n", seed);
  srand48 (seed);

  if (n_passes == 0)
    {
      n_passes = 10;
    }

  maybe_fprintf (stderr, "using %u data passes per m-step\n", n_passes);

  if (n_threads == 0)
    {
      n_threads = 1;
    }

  maybe_fprintf (stderr,
                 "using %u thread%s\n",
                 n_threads,
                 (n_threads > 1) ? "s" : "");

  if (eta == 0)
    {
      eta = 1;
    }

  maybe_fprintf (stderr, "using eta of %f\n", eta);

  if (decay == 0)
    {
      decay = 0.707;
    }

  maybe_fprintf (stderr, "using decay of %f\n", decay);

  maybe_fprintf (stderr, "reading data ... ");

  FILE* dfp = (argc < 2 || strcmp (argv[1], "-") == 0) ? stdin
                                                       : fopen (argv[1], "r");

  if (dfp == NULL)
    {
      fprintf (stderr, 
               "fatal: can't open %s: %s\n",
               argv[1],
               strerror (errno));
      exit (1);
    }

  if (read_data (dfp, &data, rank) != 0)
    {
      fprintf (stderr, "fatal: read_data failed\n");
      exit (1);
    }

  maybe_fprintf (stderr,
                 "success: n_workers = %u, n_labels = %u, n_ratings = %u\n",
                 data.n_workers,
                 data.n_labels,
                 data.n_ratings);

  maybe_fprintf (stderr, "reading clamp data ... ");

  FILE* cfp = (argc < 3) ? dfp 
                         : strcmp (argv[2], "-") == 0 ? stdin
                                                      : fopen (argv[2], "r");

  if (cfp == NULL)
    {
      fprintf (stderr, 
               "fatal: can't open %s: %s\n",
               argv[2],
               strerror (errno));
      exit (1);
    }

  if (read_clamp (cfp, &data) == 0)
    {
      maybe_fprintf (stderr, "success: n_clamp = %u\n", data.n_clamp);
      em (&data, eta, decay, n_passes, n_threads, quiet);
      print_parameters (stdout, &data);
    }

  if (cfp != dfp)
    {
      fclose (cfp);
    }

  fclose (dfp);
  data_destruct (&data);

  return 0;
}
